#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: Kailash Budhathoki (kailash.buki@gmail.com)
"""Synthetic class labelled dataset generator (inspired from Jilles' work).
The gen.json config file contains all the needed information. Importantly,
the num_alphabets field must have larger positive integer than any of the
injected ids.
"""

import json
import random


with open('gen.json') as conf_file:
    config = json.load(conf_file)

given_alphabets = set()
generated_alphabets = set()
alphabets = set()
class_labels = []

for split_info in config['split_info']:
    for inject_info in split_info['inject']:
        given_alphabets.update(inject_info['itemset'])

for generated_alphabet in xrange(config['num_alphabets']):
    if generated_alphabet not in given_alphabets:
        generated_alphabets.add(generated_alphabet)

alphabets = given_alphabets | generated_alphabets
class_label = max(alphabets) + 1
for i in xrange(config['num_splits']):
    class_labels.append(class_label)
    class_label += 1

def get_itemsets_to_inject(split_index):
    itemsets = []
    freqs = []
    this_split_info = config['split_info'][split_index]
    for inject_info in this_split_info['inject']:
        itemsets.append(set(inject_info['itemset']))
        freqs.append(inject_info['freq'])
    return itemsets, freqs

def create_split_data(split_index):
    split_data = ''
    split_stats = '-' * 100 + '\n'
    class_label = class_labels[split_index]
    num_rows = config['split_info'][split_index]['num_rows']
    itemsets, freqs = get_itemsets_to_inject(split_index)
    supports = [0] * len(itemsets)
    counts = [0] * len(itemsets)
    rows_added = 0

    while True:
        if rows_added == num_rows:
            break

        row = set()

        for i, itemset in enumerate(itemsets):
            if random.random() < freqs[i]:
                row = row.union(itemset);
                supports[i] += 1

        for alphabet in generated_alphabets:
            if alphabet not in itemsets and random.random() < config['uniform_noise']:
                row.add(alphabet)

        for i, itemset in enumerate(itemsets):
            if len(itemset & row) == len(itemset):
                counts[i] += 1

        if len(row) == 0:
            continue

        row = str(class_label) + ' ' + ' '.join(str(alphabet) for alphabet in list(sorted(row))) + '\n'
        split_data += row
        rows_added += 1

    for i, itemset in enumerate(itemsets):
        split_stats += ' '.join(str(alphabet) for alphabet in list(itemset)) + \
                    '\t	(' + str(supports[i]) + ', ' + str(round(supports[i]/float(num_rows), 2)) + ', ' + str(freqs[i]) + ')' + \
                    '\t[' + str(counts[i]) + ', ' + str(round(counts[i]/float(num_rows), 2)) + ']\n'

    return split_data, split_stats

dataset = ""
stats = "itemset (planting count, planting freq, desired freq) [support, freq]\n"
for i in range(0, config['num_splits']):
    split_data, split_stats = create_split_data(i)
    dataset += split_data
    stats += split_stats

with open(config['output'], 'w') as output_file:
    output_file.write(dataset)

with open(config['stats'], 'w') as stats_file:
    stats_file.write(stats)
